"""
Wan2.1-T2V-14B API服务
支持文生视频推理（基于MindIE）
"""
import os
import glob
import logging
import uuid
import time
import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Optional, Dict, Any
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
import shutil
import uvicorn
from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST
from starlette.responses import Response

from .model_loader import Wan2ModelLoader
from .inference import Wan2Inference

# 配置日志
logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Prometheus指标
request_counter = Counter(
    "wan2_requests_total",
    "Total number of requests",
    ["endpoint", "status"]
)

request_duration = Histogram(
    "wan2_request_duration_seconds",
    "Request duration in seconds",
    ["endpoint"]
)

# 全局变量
model_loader = None
inference_engine = None
task_store: Dict[str, Dict[str, Any]] = {}  # 任务存储

# 专用线程池执行器（用于推理任务，避免阻塞主事件循环）
inference_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="inference_")

# 输入文件目录
INPUT_DIR = "/app/inputs"


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    global model_loader, inference_engine
    
    # 启动时加载模型
    logger.info("正在加载Wan2.1-T2V模型...")
    try:
        model_loader = Wan2ModelLoader()
        model_loader.load_model()
        inference_engine = Wan2Inference(
            model_loader,
            model_loader.get_config()
        )
        logger.info("Wan2推理引擎初始化完成")
    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        raise
    
    yield
    
    # 关闭时清理资源
    logger.info("关闭推理线程池...")
    inference_executor.shutdown(wait=False)
    logger.info("正在清理资源...")


app = FastAPI(
    title="Wan2.1-T2V API",
    description="基于Wan2.1-T2V-14B的文生视频服务（MindIE推理）",
    version="1.0.0",
    lifespan=lifespan
)

# CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# 请求和响应模型
class VideoRequest(BaseModel):
    """视频生成请求"""
    model: str = Field(default="Wan2.1-T2V-14B", description="模型名称")
    prompt: str = Field(..., description="文本提示词")
    seed: Optional[int] = Field(None, description="随机种子")


class VideoSubmitResponse(BaseModel):
    """视频提交响应"""
    request_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态")


class VideoStatusResponse(BaseModel):
    """视频状态响应"""
    request_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态: pending/processing/success/failed")
    video_url: Optional[str] = Field(None, description="视频URL")
    error: Optional[str] = Field(None, description="错误信息")
    created_at: str = Field(..., description="创建时间")
    completed_at: Optional[str] = Field(None, description="完成时间")


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    model_loaded: bool
    model_id: str
    version: str


@app.get("/", tags=["Info"])
async def root():
    """根路径"""
    return {
        "service": "Wan2.1-T2V API",
        "model": "Wan2.1-T2V-14B",
        "version": "1.0.0",
        "description": "文生视频推理服务（基于MindIE）",
        "endpoints": {
            "submit": "/v1/video/submit",
            "status": "/v1/video/status/{request_id}",
            "download": "/v1/video/download/{request_id}",
            "health": "/health",
            "metrics": "/metrics"
        }
    }


@app.get("/health", response_model=HealthResponse, tags=["Info"])
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "model_loaded": model_loader is not None and inference_engine is not None,
        "model_id": model_loader.model_id if model_loader else "unknown",
        "version": "1.0.0"
    }


@app.get("/metrics", tags=["Info"])
async def metrics():
    """Prometheus指标"""
    return Response(content=generate_latest(), media_type=CONTENT_TYPE_LATEST)


async def process_video_task(request_id: str, prompt: str, seed: Optional[int], image_path: Optional[str] = None):
    """后台处理视频生成任务"""
    try:
        # 更新状态为处理中
        task_store[request_id]["status"] = "processing"
        
        logger.info(f"开始处理任务 {request_id}")
        if image_path:
            logger.info(f"输入图片: {image_path}")
        
        # 在专用线程池中异步执行推理（避免阻塞主事件循环）
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            inference_executor,
            lambda: inference_engine.generate_video(
                prompt=prompt,
                num_inference_steps=None,
                seed=seed,
                request_id=request_id,
                image_path=image_path
            )
        )
        
        # 更新任务状态
        task_store[request_id].update({
            "status": "success",
            "video_path": result["video_path"],
            "video_dir": result.get("video_dir"),
            "video_url": f"/v1/video/download/{request_id}",
            "completed_at": datetime.now().isoformat(),
            "result": result
        })
        
        logger.info(f"任务 {request_id} 完成")
        
    except Exception as e:
        logger.error(f"任务 {request_id} 失败: {str(e)}", exc_info=True)
        task_store[request_id].update({
            "status": "failed",
            "error": str(e),
            "completed_at": datetime.now().isoformat()
        })


@app.post("/v1/video/submit", response_model=VideoSubmitResponse, tags=["Video"])
async def submit_video_generation(
    prompt: str = Form(..., description="文本提示词"),
    seed: Optional[int] = Form(None, description="随机种子"),
    image: Optional[UploadFile] = File(None, description="输入图片（可选）")
):
    """提交视频生成任务（支持图片上传）"""
    endpoint = "/v1/video/submit"
    request_counter.labels(endpoint=endpoint, status="received").inc()
    
    try:
        # 生成任务ID
        request_id = str(uuid.uuid4()).replace('-', '') + str(int(time.time() * 1000))
        
        # 处理图片上传
        image_path = None
        if image and image.filename:
            # 创建输入目录: /app/inputs/$request_id
            input_dir = os.path.join(INPUT_DIR, request_id)
            os.makedirs(input_dir, exist_ok=True)
            
            # 保存图片文件
            file_ext = os.path.splitext(image.filename)[1] or '.jpg'
            image_filename = f"input{file_ext}"
            image_path = os.path.join(input_dir, image_filename)
            
            with open(image_path, "wb") as f:
                content = await image.read()
                f.write(content)
            
            logger.info(f"图片已保存: {image_path}")
        
        # 初始化任务状态
        task_store[request_id] = {
            "request_id": request_id,
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "request": {"prompt": prompt, "seed": seed, "has_image": image_path is not None}
        }
        
        # 使用 asyncio.create_task 异步执行（立即返回响应）
        asyncio.create_task(process_video_task(request_id, prompt, seed, image_path))
        
        request_counter.labels(endpoint=endpoint, status="success").inc()
        
        logger.info(f"任务 {request_id} 已提交")
        
        return VideoSubmitResponse(
            request_id=request_id,
            status="pending"
        )
        
    except Exception as e:
        logger.error(f"提交任务失败: {str(e)}", exc_info=True)
        request_counter.labels(endpoint=endpoint, status="error").inc()
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/video/status/{request_id}", response_model=VideoStatusResponse, tags=["Video"])
async def get_video_status(request_id: str):
    """查询视频生成任务状态"""
    endpoint = "/v1/video/status"
    request_counter.labels(endpoint=endpoint, status="received").inc()
    
    try:
        if request_id not in task_store:
            raise HTTPException(status_code=404, detail="任务不存在")
        
        task = task_store[request_id]
        
        request_counter.labels(endpoint=endpoint, status="success").inc()
        
        return VideoStatusResponse(
            request_id=task["request_id"],
            status=task["status"],
            video_url=task.get("video_url"),
            error=task.get("error"),
            created_at=task["created_at"],
            completed_at=task.get("completed_at")
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"查询任务状态失败: {str(e)}", exc_info=True)
        request_counter.labels(endpoint=endpoint, status="error").inc()
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/video/download/{request_id}", tags=["Video"])
async def download_video(request_id: str):
    """下载生成的视频"""
    try:
        if request_id not in task_store:
            raise HTTPException(status_code=404, detail="任务不存在")
        
        task = task_store[request_id]
        
        if task["status"] != "success":
            raise HTTPException(
                status_code=400, 
                detail=f"视频未就绪，当前状态: {task['status']}"
            )
        
        video_path = task.get("video_path")
        
        # 如果直接路径存在，使用它
        if video_path and os.path.exists(video_path):
            return FileResponse(
                video_path,
                media_type="video/mp4",
                filename=f"{request_id}.mp4"
            )
        
        # 否则在任务目录下查找mp4文件
        video_dir = task.get("video_dir")
        if video_dir and os.path.exists(video_dir):
            mp4_files = glob.glob(os.path.join(video_dir, "*.mp4"))
            if mp4_files:
                return FileResponse(
                    mp4_files[0],
                    media_type="video/mp4",
                    filename=f"{request_id}.mp4"
                )
        
        raise HTTPException(status_code=404, detail="视频文件不存在")
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"下载视频失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    port = int(os.getenv("PORT", 8000))
    
    uvicorn.run(
        "api.main:app",
        host="0.0.0.0",
        port=port,
        workers=1,
        log_level=os.getenv("LOG_LEVEL", "info").lower()
    )
